package demo.db.index;

import java.util.ArrayList;
import java.util.List;

/**
 * B+树
 */
public class BPlusTree {
    BPlusTreeNode root;
    int t;  // 最小度数

    public BPlusTree(int t) {
        this.t = t;
        root = new BPlusTreeNode(t, true);
    }

    public Object search(int key) {
        return root.search(key);
    }

    public void insert(int key, Object value) {
        if (root.keyCount == 2 * t) {
            // 根节点已满，需要分裂
            BPlusTreeNode s = new BPlusTreeNode(t, false);
            s.children[0] = root;
            s.splitChild(0, root);

            // 决定哪个子节点将接收新关键字
            int i = 0;
            if (s.keys[0] < key) {
                i++;
            }
            s.children[i].insert(key, value);

            root = s;
        } else {
            root.insert(key, value);
        }
    }

    // 范围查询(利用叶子节点链表)
    public List<Object> rangeQuery(int start, int end) {
        List<Object> result = new ArrayList<>();
        BPlusTreeNode node = findLeaf(start);

        while (node != null) {
            for (int i = 0; i < node.keyCount; i++) {
                if (node.keys[i] >= start && node.keys[i] <= end) {
                    result.add(node.values[i]);
                }
                if (node.keys[i] > end) {
                    return result;
                }
            }
            node = node.next;
        }
        return result;
    }

    // 辅助方法：找到包含key的叶子节点
    private BPlusTreeNode findLeaf(int key) {
        BPlusTreeNode node = root;
        while (!node.isLeaf) {
            int i = 0;
            while (i < node.keyCount && key >= node.keys[i]) {
                i++;
            }
            node = node.children[i];
        }
        return node;
    }
}


class BPlusTreeNode {
    int[] keys;             // 关键字数组
    int t;                  // 最小度数
    BPlusTreeNode[] children;// 子节点指针
    BPlusTreeNode next;      // 指向下一个叶子节点(仅叶子节点使用)
    Object[] values;        // 数据指针/值(仅叶子节点使用)
    int keyCount;           // 当前关键字数量
    boolean isLeaf;         // 是否为叶子节点

    public BPlusTreeNode(int t, boolean isLeaf) {
        this.t = t;
        this.isLeaf = isLeaf;
        keys = new int[2 * t];
        children = new BPlusTreeNode[2 * t + 1];
        values = (isLeaf) ? new Object[2 * t] : null;
        keyCount = 0;
        next = null;
    }

    // 查找关键字
    public Object search(int key) {
        int i = 0;
        while (i < keyCount && key >= keys[i]) {
            i++;
        }

        if (isLeaf) {
            if (i > 0 && keys[i - 1] == key) {
                return values[i - 1];  // 找到
            }
            return null;  // 未找到
        }

        return children[i].search(key); // 递归查找子节点
    }

    // 插入辅助方法
    public void insert(int key, Object value) {
        int i = keyCount - 1;

        if (isLeaf) {
            // 叶子节点插入
            while (i >= 0 && keys[i] > key) {
                keys[i + 1] = keys[i];
                values[i + 1] = values[i];
                i--;
            }
            keys[i + 1] = key;
            values[i + 1] = value;
            keyCount++;
        } else {
            // 内部节点插入
            while (i >= 0 && keys[i] > key) {
                i--;
            }

            if (children[i + 1].keyCount == 2 * t) {
                splitChild(i + 1, children[i + 1]);
                if (keys[i + 1] < key) {
                    i++;
                }
            }
            children[i + 1].insert(key, value);
        }
    }

    // 分裂子节点
    public void splitChild(int i, BPlusTreeNode y) {
        BPlusTreeNode z = new BPlusTreeNode(y.t, y.isLeaf);
        z.keyCount = t;

        if (y.isLeaf) {
            // 叶子节点分裂
            for (int j = 0; j < t; j++) {
                z.keys[j] = y.keys[j + t];
                z.values[j] = y.values[j + t];
            }
            y.keyCount = t;
            z.next = y.next;
            y.next = z;
        } else {
            // 内部节点分裂
            for (int j = 0; j < t - 1; j++) {
                z.keys[j] = y.keys[j + t + 1];
            }
            for (int j = 0; j < t; j++) {
                z.children[j] = y.children[j + t + 1];
            }
            y.keyCount = t;
        }

        // 调整当前节点
        for (int j = keyCount; j > i; j--) {
            children[j + 1] = children[j];
        }
        children[i + 1] = z;

        for (int j = keyCount - 1; j >= i; j--) {
            keys[j + 1] = keys[j];
        }
        keys[i] = y.keys[t];
        keyCount++;
    }
}